import sys
import os
import pandas as pd
import numpy as np

# 将项目根目录添加到Python路径中，以便导入dragonquant
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import dragonquant as dq
from dragonquant.api import order, log

# --- 技术指标辅助函数 ---
# 在一个成熟的库中，这些函数应该位于 dragonquant.indicators 模块
def calculate_sma(series, window):
    return series.rolling(window=window).mean()

def calculate_rsi(series, window=14):
    delta = series.diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=window).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=window).mean()
    rs = gain / loss
    return 100 - (100 / (1 + rs))

def calculate_atr(high, low, close, window=14):
    high_low = high - low
    high_close = np.abs(high - close.shift())
    low_close = np.abs(low - close.shift())
    tr = pd.DataFrame({'hl': high_low, 'hc': high_close, 'lc': low_close}).max(axis=1)
    return tr.rolling(window=window).mean()

# --- 策略实现 ---

class AtrRsiMaStrategy(dq.Strategy):
    """
    一个系统的中期回撤买入策略：
    1. 趋势过滤：只在SMA(120)之上寻找买入机会。
    2. 入场信号：当RSI(14)回调至低位(如35)时买入。
    3. 风险管理：使用ATR(14)设置动态止损位。
    4. 出场信号：趋势反转、RSI超买或触及止损时卖出。
    """
    def init(self):
        # 1. 设置策略参数
        self.stock = '510500.SH'  # 交易标的：中证500 ETF，流动性好，代表性强
        self.sma_long_window = 120
        self.rsi_window = 14
        self.rsi_oversold = 35
        self.rsi_overbought = 75
        self.atr_window = 14
        self.atr_multiplier = 2.0
        
        # 2. 初始化策略状态变量
        self.stop_loss_price = 0.0 # 用于跟踪当前持仓的止损价
        self.in_market = False     # 标记当前是否持仓

        log(f"策略 {self.__class__.__name__} 初始化完成。交易标的: {self.stock}")

    def next(self, context):
        # 1. 获取所需的最长历史数据
        required_bars = self.sma_long_window + 5 # 加一点buffer
        hist = context.data.history(self.stock, ['open', 'high', 'low', 'close'], required_bars, context.dt)
        
        if len(hist) < required_bars:
            return # 数据不足，跳过

        current_price = hist['close'].iloc[-1]
        
        # 2. 计算所有需要的技术指标
        sma_long = calculate_sma(hist['close'], self.sma_long_window).iloc[-1]
        rsi = calculate_rsi(hist['close'], self.rsi_window).iloc[-1]
        atr = calculate_atr(hist['high'], hist['low'], hist['close'], self.atr_window).iloc[-1]

        # 3. 根据是否持仓，执行不同的逻辑
        if not self.in_market:
            # --- 入场逻辑 ---
            is_uptrend = current_price > sma_long
            is_oversold_pullback = rsi < self.rsi_oversold

            if is_uptrend and is_oversold_pullback:
                # 计算买入数量，使用95%的现金
                cash_to_use = context.portfolio.cash * 0.95
                shares_to_buy = int(cash_to_use / current_price)
                if shares_to_buy > 0:
                    order(self.stock, shares_to_buy)
                    
                    # 设置止损位并更新状态
                    self.stop_loss_price = current_price - self.atr_multiplier * atr
                    self.in_market = True
                    log(f"[{context.dt.date()}] 买入信号触发: 价格={current_price:.2f}, RSI={rsi:.2f}."
                        f" 买入 {shares_to_buy} 股. 止损位设为: {self.stop_loss_price:.2f}")
        else:
            # --- 出场逻辑 ---
            # 检查三个出场条件
            stop_loss_triggered = current_price <= self.stop_loss_price
            trend_reversed = current_price < sma_long
            take_profit_triggered = rsi > self.rsi_overbought

            if stop_loss_triggered:
                reason = f"止损触发 (价格 <= {self.stop_loss_price:.2f})"
            elif trend_reversed:
                reason = f"趋势反转 (价格 < SMA{self.sma_long_window})"
            elif take_profit_triggered:
                reason = f"止盈触发 (RSI > {self.rsi_overbought})"
            else:
                return # 无出场信号，继续持有

            # 执行卖出
            order(self.stock, 0, style='target') # 'target'风格表示卖出至目标仓位0，即清仓
            log(f"[{context.dt.date()}] 卖出信号触发: {reason}. 价格={current_price:.2f}, RSI={rsi:.2f}. 清仓。")
            
            # 重置状态变量
            self.stop_loss_price = 0.0
            self.in_market = False


if __name__ == '__main__':
    # 1. 初始化回测引擎
    cerebro = dq.Cerebro()

    # 2. 添加策略
    cerebro.add_strategy(AtrRsiMaStrategy)

    # 3. 设置股票池（必须）
    cerebro.set_universe(['510500.SH'])
    # 4. 设置回测参数
    cerebro.configure(
        start_date='2020-01-01',
        end_date='2024-12-31',
        initial_cash=1000000,
        data_source='real',
        data_path='./data',
        benchmark='000300.SH'
    )
    
    # 在真实的库中，这里会设置交易费用
    # cerebro.set_commission(commission=0.0003, stamp_duty=0.001)

    # 4. 运行回测
    # 这里的context, data等对象都是由cerebro在内部创建和管理的
    # 为了演示，我们假设run()方法会模拟这个过程
    print("假设`run()`方法正在模拟回测过程...")
    # 伪代码: 
    # context.portfolio = Portfolio(initial_cash)
    # for date in trading_days:
    #     context.dt = date
    #     context.data = DataHandler(...)
    #     strategy.next(context)
    
    results = cerebro.run()

    # 5. 分析和可视化结果
    print("\n回测完成，生成分析图表...")
    results.plot()